package com.xl.bigdata.es.index;

import com.xl.bigdata.es.xpack.ESXpackRestHighLevelClient;
import org.apache.log4j.Logger;
import org.apache.log4j.spi.LoggerFactory;
import org.apache.logging.log4j.spi.LoggerContext;
import org.elasticsearch.action.bulk.*;
import org.elasticsearch.action.index.IndexRequest;
import org.elasticsearch.action.support.WriteRequest;
import org.elasticsearch.client.RequestOptions;
import org.elasticsearch.client.RestHighLevelClient;
import org.elasticsearch.common.unit.ByteSizeUnit;
import org.elasticsearch.common.unit.ByteSizeValue;
import org.elasticsearch.common.xcontent.XContentType;
import org.elasticsearch.core.TimeValue;

import java.io.IOException;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.UUID;
import java.util.concurrent.TimeUnit;

/**
 * @program: IAEngine
 * @description: 批量写入数据
 * @author: XL.Gao
 * @create: 2021-08-30 10:24
 **/
public class BulkIndexAPI {


    /**
     * @param client:
     * @param map:
     * @param index:
     * @param onlyId: true，使用Map的key作为主键
     * @Description:
     * @Author: xl.gao
     * @Date:
     * @return: boolean
     **/
    public static boolean bulkRequestIndexAPI(RestHighLevelClient client, Map<String, String> map, String index, boolean onlyId, boolean refresh) {
        if (map == null || map.size() == 0) {
            return false;
        }

        BulkRequest bulkRequest = new BulkRequest();

        for (String key : map.keySet()) {
            IndexRequest indexRequest = new IndexRequest(index);
            if (onlyId) {
                indexRequest.id(key);
            }
            indexRequest.source(map.get(key), XContentType.JSON);

            bulkRequest.add(indexRequest);

        }

        if (refresh) {
            // 刷新，保证索引可见
            bulkRequest.setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE);
        }

        try {

            // bulkRequest.numberOfActions() 必须重新new才会清零
            if (bulkRequest.numberOfActions() > 0) {

                System.out.println("numberOfActions:" + bulkRequest.numberOfActions());

                client.bulk(bulkRequest, RequestOptions.DEFAULT);

//                BulkResponse bulkResponse = client.bulk(bulkRequest, RequestOptions.DEFAULT);
//
//                // 如果有失败，返回true
//                if (bulkResponse.hasFailures()) {
//                    for (BulkItemResponse bulkItemResponse : bulkResponse) {
//                        if(bulkItemResponse.isFailed()){
//                            BulkItemResponse.Failure bulkItemResponseFailure = bulkItemResponse.getFailure();
//                            System.out.println("Error -- " + bulkItemResponseFailure.getCause());
//                        }
//                    }
//
//                    return false;
//                }
            }

            bulkRequest = null;

        } catch (IOException e) {
            e.printStackTrace();
        }

        return true;
    }

    public static boolean bulkProcessorBuilderIndexAPI(RestHighLevelClient client, Map<String, String> map, String index, boolean onlyId, boolean refresh) {

        BulkProcessor.Listener listener = new BulkProcessor.Listener() {
            @Override
            public void beforeBulk(long executionId, BulkRequest request) {
                int numberOfActions = request.numberOfActions();

                System.out.println("Executing bulk " + executionId + " with " + numberOfActions + " requests");

            }

            @Override
            public void afterBulk(long executionId, BulkRequest request,
                                  BulkResponse response) {
                if (response.hasFailures()) {
                    System.out.println("Bulk " + executionId + " executed with failures");
                } else {
                    System.out.println("Bulk " + executionId + " completed in " + response.getTook().getMillis() + " milliseconds");
                }
            }

            @Override
            public void afterBulk(long executionId, BulkRequest request,
                                  Throwable failure) {
                System.out.println("Failed to execute bulk" + failure);
            }
        };

        BulkProcessor.Builder builder = BulkProcessor.builder(
                (request, bulkListener) ->
                        client.bulkAsync(request, RequestOptions.DEFAULT, bulkListener),
                listener, "--builder");
        builder.setBulkActions(1000);
        builder.setBulkSize(new ByteSizeValue(8L, ByteSizeUnit.MB));
        builder.setConcurrentRequests(1);
        builder.setFlushInterval(TimeValue.timeValueSeconds(20L));
        builder.setBackoffPolicy(BackoffPolicy
                .constantBackoff(TimeValue.timeValueSeconds(1L), 3));

        BulkProcessor buildProcessor = builder.build();
        for (String key : map.keySet()) {
            IndexRequest indexRequest = new IndexRequest(index);
            if (onlyId) {
                indexRequest.id(key);
            }
            indexRequest.source(map.get(key), XContentType.JSON);
            buildProcessor.add(indexRequest);

        }

        try {
            boolean terminated = buildProcessor.awaitClose(30L, TimeUnit.SECONDS);
        } catch (InterruptedException e) {
            e.printStackTrace();
        }
        buildProcessor.close();
        return true;
    }


    public static void main(String[] args) {

        String index = "address_base";
        String id = "HCC200418WR9338B";

        RestHighLevelClient restHighLevelClient = ESXpackRestHighLevelClient.getRestHighLevelClient();

        String api = GetAPI.getAPI(restHighLevelClient, index, id);

        Map<String, String> map = new HashMap<>();


        long l = System.currentTimeMillis();
        for (int i = 0; i < 1000000; i++) {
            map.put(UUID.randomUUID().toString(), api);

            if (i % 300000 == 0) {
                // 10000 93655
//                BulkIndexAPI.bulkRequestIndexAPI(restHighLevelClient, map, index, false, true);
                // 10000 68196
                bulkProcessorBuilderIndexAPI(restHighLevelClient, map, index, false, true);
                System.out.println(i);
                map = new HashMap<>();
            }
        }
//        // 100 0000 ,69529
//        bulkProcessorBuilderIndexAPI(restHighLevelClient, map, index, false, true);

        System.out.println(System.currentTimeMillis() - l);

        ESXpackRestHighLevelClient.closeRestHighLevelClient();

    }
}
